module(..., package.seeall)


multiset = class("multiset")
_G.multiset = multiset

local weight_mt = {__index=function(t, k) return 0 end}

function multiset.__init(self, weights)
    self.weight = setmetatable({}, weight_mt)
    self.size = 0
    if weights ~= nil then
        for elem, weight in pairs(weights) do
            self:add(elem, weight)
        end
    end
end

function multiset.__tostring(self)
    return self.__name .. table.repr(self.weight)
end

function multiset.__len(self)
    return self.size
end

function multiset.__iter(self)
    return next, self.weight, nil
end
multiset.iter = multiset.__iter

function multiset.contains(self, elem)
    return self.weight[elem] > 0
end

function multiset.add(self, elem, weight)
    local weight = weight or 1
    local new_weight = self.weight[elem] + weight
    if new_weight < 0 then
        error("provided arguments make weight negative in add()", 2)
    elseif new_weight > 0 then
        self.weight[elem] = new_weight
    else
        self.weight[elem] = nil
    end
    self.size = self.size + weight
    return new_weight
end

function multiset.remove(self, elem, weight)
    --[[ Removing an element from a multiset is the same as adding a negative weight to it. If no 
    weight is specified, the element is completely removed from the multiset/bag. ]]
    weight = weight or self.weight[elem]
    if weight <= 0 then
        error("positive weight required in remove()", 2)
    end
    return self:add(elem, -weight)
end

function multiset.discard(self, elem)
    if self.weight[elem] > 0 then
        self:remove(elem)
    end
end

function multiset.pop(self, elem)
    --[[ Pops a specific element from the multiset, returning the element and its weight prior to 
    removal. If no element is specified, an arbitrary element is popped. ]]
    if self.size == 0 then
        error("cannot pop() from empty multiset", 2)
    end
    if elem == nil then
        elem = next(self.weight)
    end
    local weight = self.weight[elem]
    self:remove(elem)
    return elem, weight
end

function multiset.clear(self)
    self.weight = setmetatable({}, weight_mt)
    self.size = 0
end


function test()
    local s = multiset{a=1, b=2, c=3, d=4}
    print(s)
    assert(s.weight["a"] == 1)
    assert(s.weight["b"] == 2)
    assert(s.weight["c"] == 3)
    assert(s.weight["d"] == 4)
    assert(s.weight["foobar"] == 0)
    
    s:add("a")
    print(s)
    assert(s.weight["a"] == 2)
    
    s:remove("a")
    print(s)
    assert(s.weight["a"] == 0)
    
    local b, b_weight = s:pop("b")
    assert(b_weight == 2)
    print(s)
    assert(s.weight["b"] == 0)
    assert(s:contains("c"))
    assert(s.size == 7)
    return s
end
